function [v, c, savings, A] = solveHJB(r, vInit, incomeStruct, prefStruct, Xstruct, simStruct);
%{
%%-------------------------------------------------------------------------
Solves the HJB of the 2-state Huggett economy given interest rate r

    If beta < 1, this code solves the Bellman of the u-hat agent
                 then backs out IG agent's behavior
%%-------------------------------------------------------------------------
%}

%%-------------------------------------------------------------------------
%Unpack Structs / Setup
    
    lambda_L = incomeStruct.lambda_L;
    lambda_H = incomeStruct.lambda_H;
    y_L = incomeStruct.y_L;
    y_H = incomeStruct.y_H;
        y_ = [y_L, y_H];
        ybar = (lambda_H*y_L + lambda_L*y_H)/(lambda_L + lambda_H);

    rho = prefStruct.rho;
    beta = prefStruct.beta;
    betaE = prefStruct.betaE;
    gamma = prefStruct.gamma;
    psi = prefStruct.psi;
    util = prefStruct.util;
    util_hat = prefStruct.util_hat;

    borrowingLimit = Xstruct.borrowingLimit;
    xmax = Xstruct.xmax;
    xjump = Xstruct.xjump;
    I = Xstruct.I;
    X_ = Xstruct.X_;
    r_ = r*ones(I, 1);
    
    maxIterations = simStruct.maxIterations;
	crit = simStruct.crit;
	Delta = simStruct.Delta;

    X__ = [X_, X_];
    y__ = ones(I, 1)*y_;
    
%%-------------------------------------------------------------------------

%%-------------------------------------------------------------------------
%Solve HJB
    dVf = zeros(I, 2);
    dVb = zeros(I, 2);
    c = zeros(I, 2);
    savings = zeros(I,2);
    c_shifted = zeros(I, 2);

    Aswitch = [-speye(I)*lambda_L, speye(I)*lambda_L; speye(I)*lambda_H, -speye(I)*lambda_H];

    %Initial Guess
    if vInit == 0
        v0(:,1) = util_hat(ybar + max(r, 0.005).*X_)./ rho;
        v0(:,2) = util_hat(ybar + max(r, 0.005).*X_)./ rho;
        solveNoUboost = 1; %To get good initial starting point, don't allow for "utility boost" at boundary
    else
        v0 = vInit;
        solveNoUboost = 0;
    end
    v = v0;

    for n=1:maxIterations
        V = v;
        Vstore___(:,:,n) = V;
        
        % Forward Difference
        dVf(1:I-1,:) = (V(2:I,:)-V(1:I-1,:))/xjump;
        dVf(I,:) = ((psi^gamma)/betaE)*((y_ + r_(I).*xmax).^(-gamma));
        
        % Backward Difference
        dVb(2:I,:) = (V(2:I,:) - V(1:I-1,:))/xjump;
        %lower boundary condition imposed directly on consumption below
		
        %consumption and savings with forward difference
        cf = max((betaE/(psi^gamma))*dVf, 10^(-10)).^(-1/gamma);
            %Impose boundary: must move forward if not stopping
            cf(1,:) = min(cf(1,:), psi*(y_ + r_(1).*borrowingLimit));
            dVf(1,:) = ((psi^gamma)/betaE)*(cf(1,:).^(-gamma));
        ssf = y__ + r_.*X__ - cf;
        Hf = util_hat(cf) + dVf.*ssf;
        
        %consumption and savings with backward difference
        cb = max((betaE/(psi^gamma))*dVb, 10^(-10)).^(-1/gamma);
        ssb = round(y__ + r_.*X__ - cb, 15);
            ssb(1,:) = 0;
        Hb = util_hat(cb) + dVb.*ssb;
        
        %consumption at steady state
        c0 = y__ + r_.*X__;
        H0 = util_hat(c0);

        % dV_upwind makes choice of forward/backward based on sign of drift    
        Ineither = (1-(ssf>0)) .* (1-(ssb<0));
        Iunique = (ssb<0).*(1-(ssf>0)) + (1-(ssb<0)).*(ssf>0);
        Iboth = (ssb<0).*(ssf>0);
        Ib = Iunique.*(ssb<0) + Iboth.*(Hb>Hf);
        If = Iunique.*(ssf>0) + Iboth.*(Hf>=Hb);
        I0 = Ineither;

        c  = cf.*If + cb.*Ib + c0.*I0;
        u = util_hat(c);
            
        %--------------------------------------------------------------

        %CONSTRUCT MATRIX
        X = -Ib.*ssb / xjump;
        Y = -If.*ssf / xjump + Ib.*ssb / xjump;
        Z = If.*ssf / xjump;

        A1=spdiags(Y(:,1),0,I,I)+spdiags(X(2:I,1),-1,I,I)+spdiags([0;Z(1:I-1,1)],1,I,I);
        A2=spdiags(Y(:,2),0,I,I)+spdiags(X(2:I,2),-1,I,I)+spdiags([0;Z(1:I-1,2)],1,I,I);
        A = [A1,sparse(I,I);sparse(I,I),A2] + Aswitch;
        
        u_stacked = [u(:,1); u(:,2)];
        V_stacked = [V(:,1); V(:,2)];
        
        %Allow for "boosted" utility at boundary b_
        VstarL = (util(y_L+r_(1)*borrowingLimit) + lambda_L*v(1,2))./(rho + lambda_L);
        VstarH = (util(y_H+r_(1)*borrowingLimit) + lambda_H*v(1,1))./(rho + lambda_H);
        Vstar_stacked = [VstarL.*ones(I,1);VstarH.*ones(I,1)];
                
        %Set up as option-value problem (using LCP to solve)
        B = (rho + 1/Delta)*speye(2*I) - A;
        vec = u_stacked + V_stacked/Delta;
        q = -vec + B*Vstar_stacked;
        
        if solveNoUboost == 0
            %using Yuval Tassa's Newton-based LCP solver, download from http://www.mathworks.com/matlabcentral/fileexchange/20952
            z0 = V_stacked-Vstar_stacked; lbnd = zeros(2*I,1); ubnd = Inf*ones(2*I,1);
            z = LCP(B,q,lbnd,ubnd,z0,0);
            V_stacked = z+Vstar_stacked;
            
            LCP_error = max(abs(z.*(B*z + q)));
            if LCP_error > 10^(-5)
                disp('LCP not solved, Iteration =')
                disp(n)
            end
        elseif solveNoUboost == 1
            V_stacked = B\vec;
        end
        
        V = [V_stacked(1:I),V_stacked(I+1:2*I)];
        Vchange = V - v;
        v = V;

        dist(n) = max(max(abs(Vchange)));
        if dist(n) < crit && solveNoUboost == 1
            %Have initial guess, now introduce the VI element
            solveNoUboost = 0;
            continue
        end
        if dist(n) < crit && solveNoUboost == 0
            disp('Value Function Converged, Iteration = ')
            disp(n)
            break
        end
    end
    
	%--------------------------------------------------------------
    % IG Additional Steps
	% We've now solved for the u-hat agent
	%	Here we apply changes to consumption, savings, and A
	%	in order to back out IG behavior from the u-hat agent
    
	%Get c of IG from c of u-hat
    cshiftFactor = ((beta/betaE)^(-1/gamma))*(1/psi);
    c_shifted = cshiftFactor*c;
        c_shifted(1,:) = min(c_shifted(1,:), y_ + r_(1).*borrowingLimit);
    savings = (r_.*X__ + y__ - c_shifted);
    
    %Apply updated c choice to transition matrix
    ssb_shifted = round(y__ + r_.*X__ - cshiftFactor*cb, 15);
        ssb_shifted(1, :) = 0;
    ssf_shifted = y__ + r_.*X__ - cshiftFactor*cf;
        Ineither = (1-(ssf_shifted>0)) .* (1-(ssb_shifted<0));
        Iunique = (ssb_shifted<0).*(1-(ssf_shifted>0)) + (1-(ssb_shifted<0)).*(ssf_shifted>0);
        Iboth = (ssb_shifted<0).*(ssf_shifted>0);
        Ib = Iunique.*(ssb_shifted<0) + Iboth.*(Hb>=Hf);
        If = Iunique.*(ssf_shifted>0) + Iboth.*(Hf>=Hb);
        I0 = Ineither;
    X = -Ib.*ssb_shifted / xjump;
    Y = -If.*ssf_shifted / xjump + Ib.*ssb_shifted / xjump;
    Z = If.*ssf_shifted / xjump;

    A1_shifted=spdiags(Y(:,1),0,I,I)+spdiags(X(2:I,1),-1,I,I)+spdiags([0;Z(1:I-1,1)],1,I,I);
    A2_shifted=spdiags(Y(:,2),0,I,I)+spdiags(X(2:I,2),-1,I,I)+spdiags([0;Z(1:I-1,2)],1,I,I);
    A_shifted = [A1_shifted,sparse(I,I);sparse(I,I),A2_shifted] + Aswitch;
    
    %Update what gets returned
    A = A_shifted;
    c = c_shifted;
end
%%-------------------------------------------------------------------------
